from nsga2.nsga2 import NSGA2
from problems.crashworthiness import CrashworthinessDesign
import numpy as np
import matplotlib.pyplot as plt


def run_cwd():
    optimizer = NSGA2(pop_size=200, n_var=5, pc=0.9, pm=0.2, eta_c=10, eta_m=10)
    problem = CrashworthinessDesign()

    [optimum_x, optimum_fx] = optimizer.run(problem, max_gen=1000, seed=42)

    optimum_fx_array = np.array(optimum_fx)
    f1 = optimum_fx_array[:, 0]
    f2 = optimum_fx_array[:, 1]
    f3 = optimum_fx_array[:, 2]

    fig = plt.figure(figsize=(10, 7))  
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(f2, f1, f3, color='grey', s=40, alpha=0.8, label='NSGA-II Solutions')
    ax.set_xlabel('f2')
    ax.set_ylabel('f1')
    ax.set_zlabel('f3')
    ax.invert_yaxis()
    ax.legend()
    plt.title('NSGA-II Optimization Results for Crashworthiness Design')
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    run_cwd()